"""
Darknet53代码
"""
import paddle
import paddle.nn.functional as F
import numpy as np

class ConvBNLayer(paddle.nn.Layer):
    def __init__(self,ch_in,ch_out,kernel_size=3,stride=1,groups=1,padding=0,act='leaky'):
        super(ConvBNLayer,self).__init__()

        self.conv = paddle.nn.Conv2D(
            in_channels=ch_in,
            out_channels=ch_out,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            groups=groups,
            weight_attr=paddle.ParamAttr(
                initializer=paddle.nn.initializer.Normal(0.,0.02)),
            bias_attr=False)

        self.batch_norm = paddle.nn.BatchNorm2D(
            num_features=ch_out,
            weight_attr=paddle.ParamAttr(
                initializer=paddle.nn.initializer.Normal(0.,0.02),
                regularizer=paddle.regularizer.L2Decay(0.)),
            bias_attr=paddle.ParamAttr(
                initializer=paddle.nn.initializer.Constant(0.0),
                regularizer=paddle.regularizer.L2Decay(0.)))

        self.act = act


    def forward(self,inputs):
        out = self.conv(inputs)
        out = self.batch_norm(out)
        if self.act == 'leaky':
            out = F.leaky_relu(x=out,negative_slope=0.1)

        return out

class DownSample(paddle.nn.Layer):
    # 下采样，图片减半，具体实现方式是使用stride=2的卷积
    def __init__(self,
                 ch_in,
                 ch_out,
                 kernel_size=3,
                 stride=2,
                 padding=1):
        super(DownSample,self).__init__()

        self.conv_bn_layer = ConvBNLayer(
            ch_in=ch_in,
            ch_out=ch_out,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding)

        self.ch_out = ch_out

    def forward(self,inputs):
        out = self.conv_bn_layer(inputs)
        return out


class BasicBlock(paddle.nn.Layer):
    """
    基本残差块的定义，输入x经过两层卷积，然后接第二层卷积的输出和输入x相加
    """
    def __init__(self,ch_in, ch_out):
        super(BasicBlock,self).__init__()

        self.conv1 = ConvBNLayer(
            ch_in=ch_in,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0
        )

        self.conv2 = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1
        )

    def forward(self,inputs):
        conv1 = self.conv1(inputs)
        conv2 = self.conv2(conv1)
        out = paddle.add(x=inputs,y=conv2)
        return out


class LayerWarp(paddle.nn.Layer):
    """
    添加多层残差块，组成Darknet53网络的一个层级
    """
    def __init__(self,ch_in,ch_out,count,is_test=True):
        super(LayerWarp,self).__init__()

        self.basicblock0 = BasicBlock(ch_in,ch_out)
        self.res_out_list = []
        for i in range(1,count):
            # 使用add_sublayer添加子层
            res_out = self.add_sublayer("basic_block_%d" % (i),
                                        BasicBlock(ch_out*2,ch_out))
            self.res_out_list.append(res_out)

    def forward(self,inputs):
        y = self.basicblock0(inputs)
        for basic_block_i in self.res_out_list:
            y = basic_block_i(y)
        return y

# DarkNet 每组残差块的个数，来自DarkNet 的网络结构图
DarkNet_cfg = {53: ([1,2,8,8,4])}

class DarkNet53_conv_body(paddle.nn.Layer):
    def __init__(self):
        super(DarkNet53_conv_body,self).__init__()
        self.stages = DarkNet_cfg[53]
        self.stages = self.stages[0:5]

        # 第一层卷积
        self.conv0 = ConvBNLayer(
            ch_in=3,
            ch_out=32,
            kernel_size=3,
            stride=1,
            padding=1
        )

        # 下采样，使用stride=2的卷积来实现
        self.downsample0 = DownSample(ch_in=32,ch_out=32*2)

        # 添加各个层级的实现
        self.darknet53_conv_block_list = []
        self.downsample_list = []
        for i, stage in  enumerate(self.stages):
            conv_block = self.add_sublayer(
                "stage_%d" % (i),
                LayerWarp(32*(2**(i+1)),
                          32*(2**i),
                          stage))
            self.darknet53_conv_block_list.append(conv_block)

        # 两个层级之间使用DownSample将尺寸减半
        for i in range(len(self.stages) - 1):
            downsample = self.add_sublayer(
                "stage_%d_downsample" % i,
                DownSample(ch_in=32*(2**(i+1)),
                           ch_out=32*(2**(i+2))))
            self.downsample_list.append(downsample)


    def forward(self,inputs):
        out = self.conv0(inputs)
        #print("conv1:",out.numpy())
        out = self.downsample0(out)
        # print("dy:",out.numpy())
        blocks = []
        # 依次将各个层级作用在输入上面
        for i,conv_block_i in enumerate(self.darknet53_conv_block_list):
            out = conv_block_i(out)
            blocks.append(out)
            if i < len(self.stages) -1:
                out = self.downsample_list[i](out)

        return blocks[-1:-4:-1] # 将C0，C1，C2作为返回值


# backbone = DarkNet53_conv_body()
# x = np.random.randn(1,3,640,640).astype('float32')
# x = paddle.to_tensor(x)
# C0,C1,C2 = backbone(x)
# print(C0.shape,C1.shape,C2.shape)

class YoloDetectionBlock(paddle.nn.Layer):
    # define YOLOv3 detection head
    # 使用多层卷积和BN提取特征
    def __init__(self,ch_in,ch_out,is_test=True):
        super(YoloDetectionBlock,self).__init__()

        assert ch_out % 2 == 0, "channel {} cannot be divided by 2".format(ch_out)

        self.conv0 = ConvBNLayer(
            ch_in=ch_in,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0
        )
        self.conv1 = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1
        )
        self.conv2 = ConvBNLayer(
            ch_in=ch_out*2,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0
        )
        self.conv3 = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1
        )
        self.route = ConvBNLayer(
            ch_in=ch_out*2,
            ch_out=ch_out,
            kernel_size=1,
            stride=1,
            padding=0
        )
        self.tip = ConvBNLayer(
            ch_in=ch_out,
            ch_out=ch_out*2,
            kernel_size=3,
            stride=1,
            padding=1
        )

    def forward(self,inputs):
        out = self.conv0(inputs)
        out = self.conv1(out)
        out = self.conv2(out)
        out = self.conv3(out)
        route = self.route(out)
        tip = self.tip(route)
        return route,tip

# 定义Sigmoid函数
def sigmoid(x):
    return 1./(1.0 + np.exp(-x))

# 将网络特征图输出的[tx,ty,th,tw]转化成预测框的坐标[x1,y1,x2,y2]
def get_yolo_box_xxyy(pred,anchors,num_classes,downsample):
    """
    :param pred: 网络输出特征图转化为numpy.ndarray
    :param anchors: 是一个list.表示锚框的大小，列如 anchors = [116,90,156,198,373,326],表示有三个锚框
    第一个锚框大小[w,h]是[116,90],第二个锚框大小是[156,198],第三个锚框大小是[373,326]
    :param num_classes:
    :param downsample:
    :return:
    """
    batchsize = pred.shape[0]
    num_rows = pred.shape[-2]
    num_cols = pred.shape[-1]

    input_h = num_rows * downsample
    input_w = num_cols * downsample

    num_anchors = len(anchors) // 2

    # pred的形状是[N,C,H,W],其中C = NUM_ANCHORS * (5  + NUM_CLASSES)
    # 对pred进行reshape
    pred = pred.reshape([-1,num_anchors,5+num_classes,num_rows,num_cols])
    pred_location = pred[:,:,0:4,:,:]
    pred_location = np.transpose(pred_location,(0,3,4,1,2))
    anchors_this = []
    for ind in range(num_anchors):
        anchors_this.append([anchors[ind*2],anchors[ind*2+1]])
    anchors_this = np.array(anchors_this).astype('float32')

    # 最终输出数据保存在pred_box中，其中形状是[N,H,W,NUM_ANCHORS,4]
    # 其中最后一个维度4代表位置的4个坐标
    pred_box = np.zeros(pred_location.shape)
    for n in range(batchsize):
        for i in range(num_rows):
            for j in range(num_cols):
                for k in range(num_anchors):
                    pred_box[n,i,j,k,0] = j
                    pred_box[n,i,j,k,1] = i
                    pred_box[n,i,j,k,2] = anchors_this[k][0]
                    pred_box[n,i,j,k,3] = anchors_this[k][1]

    # 这里使用相对坐标，pred_box的输出元素数值在0.~1.0之间
    pred_box[:,:,:,:,0] = (sigmoid(pred_location[:,:,:,:,0]) + pred_box[:,:,:,:,0]) / num_cols
    pred_box[:,:,:,:,1] = (sigmoid(pred_location[:,:,:,:,1]) + pred_box[:,:,:,:,1]) / num_rows
    pred_box[:,:,:,:,2] = np.exp(pred_location[:,:,:,:,2]) * pred_box[:,:,:,:,2] /input_w
    pred_box[:,:,:,:,3] = np.exp(pred_location[:,:,:,:,3]) * pred_box[:,:,:,:,3] /input_h

    # 将坐标从xywh转化成xyxy
    pred_box[:,:,:,:,0] = pred_box[:,:,:,:,0] - pred_box[:,:,:,:,2] / 2.
    pred_box[:,:,:,:,1] = pred_box[:,:,:,:,1] - pred_box[:,:,:,:,3] / 2.
    pred_box[:,:,:,:,2] = pred_box[:,:,:,:,0] + pred_box[:,:,:,:,2]
    pred_box[:,:,:,:,3] = pred_box[:,:,:,:,1] + pred_box[:,:,:,:,3]

    pred_box = np.clip(pred_box,0.,1.0)

    return  pred_box


NUM_ANCHORS = 3
NUM_CLASSES = 7
num_filters = NUM_ANCHORS * (NUM_CLASSES + 5)

backbone = DarkNet53_conv_body()
detection = YoloDetectionBlock(ch_in=1024,ch_out=512)
conv2d_pred = paddle.nn.Conv2D(in_channels=1024,out_channels=num_filters,kernel_size=1)

x = np.random.randn(1,3,640,640).astype('float32')
x = paddle.to_tensor(x)
C0,C1,C2 = backbone(x)
route,tip = detection(C0)
P0 = conv2d_pred(tip)
# print(P0.shape)

reshaped_p0 = paddle.reshape(P0,[-1,NUM_ANCHORS,NUM_CLASSES+5,P0.shape[2],P0.shape[3]])
pred_objectness = reshaped_p0[:,:,4,:,:]
pred_objectness_probability = F.sigmoid(pred_objectness)
# print(pred_objectness.shape,pred_objectness_probability.shape)
pred_location = reshaped_p0[:,:,0:4,:,:]
# print(pred_location.shape)

# anchors包含了预先设定好的锚框尺寸
# anchors = [116,90,156,198,373,326]
# # downsample是特征图P0的步副
# pred_boxes = get_yolo_box_xxyy(P0.numpy(),anchors,num_classes=7,downsample=32) # 由输出特征图P0计算预测框位置坐标
# print(pred_boxes.shape)

# 取出与类别相关的预测值
pred_classification = reshaped_p0[:,:,5:5+NUM_CLASSES,:,:]
pred_classification_probability = F.sigmoid(pred_classification)
print(pred_classification.shape)

# 定义上采样模块
class Upsample(paddle.nn.Layer):
    def __init__(self, scale=2):
        super(Upsample,self).__init__()
        self.scale = scale

    def forward(self,inputs):
        # get dynamic upsample output shape
        shape_nchw = paddle.shape(inputs)
        shape_hw = paddle.slice(shape_nchw,axes=[0],starts=[2],ends=[4])
        shape_hw.stop_gradient = True
        in_shape = paddle.cast(shape_hw,dtype='int32')
        out_shape = in_shape * self.scale
        out_shape.stop_gradient = True

        # resize by actual_shape
        out = paddle.nn.functional.interpolate(
            x=inputs,scale_factor=self.scale,mode="NEAREST"
        )
        return out


class YOLOv3(paddle.nn.Layer):
    def __init__(self,num_classes=7):
        super(YOLOv3,self).__init__()

        self.num_classes = num_classes
        # 提取图像特征的骨干代码
        self.block = DarkNet53_conv_body()
        self.block_outputs = []
        self.yolo_blocks = []
        self.route_blocks_2 = []
        # 生成3个层级的特征图P0,P1,P2
        for i in range(3):
            # 添加从ci生成ri和ti的模块
            yolo_block = self.add_sublayer(
                "yolo_detecton_block_%d" % (i),
                YoloDetectionBlock(
                    ch_in = 512//(2**i)*2 if i==0 else 512//(2**i)*2 + 512//(2**i),
                    ch_out = 512//(2**i)))

            self.yolo_blocks.append(yolo_block)

            num_filters = 3 * (self.num_classes + 5)

            # 添加从ti生成pi的模块，这是一个Conv2D操作，输出通道数为3 * (num_classes + 5)
            block_out = self.add_sublayer(
                "block_out_%d" % (i),
                paddle.nn.Conv2D(in_channels=512//(2**i)*2,
                                 out_channels = num_filters,
                                 kernel_size=1,
                                 stride=1,
                                 padding=0,
                                 weight_attr=paddle.ParamAttr(
                                     initializer=paddle.nn.initializer.Normal(0.,0.02)),
                                 bias_attr=paddle.ParamAttr(
                                     initializer=paddle.nn.initializer.Constant(0.0),
                                     regularizer=paddle.regularizer.L2Decay(0.))))
            self.block_outputs.append(block_out)
            if i < 2:
                # 对ri进行卷积
                route = self.add_sublayer("route2_%d"%i,ConvBNLayer(ch_in=512//(2**i),
                                                                    ch_out=256//(2**i),
                                                                    kernel_size=1,
                                                                    stride=1,
                                                                    padding=0))

                self.route_blocks_2.append(route)

            # 将ri放大以便跟c_{i+1}保持同样的尺寸
            self.upsample = Upsample()

    def forward(self,inputs):
        outputs = []
        blocks = self.block(inputs)
        for i ,block in enumerate(blocks):
            if i > 0:
                # 将r_{i-1}经过卷积和上采样之后得到特征图，与这一级的ci进行拼接
                block = paddle.concat([route,block],axis=1)

            # 从ci生成ti和ri
            route, tip = self.yolo_blocks[i](block)
            # 从ti生成pi
            block_out = self.block_outputs[i](tip)
            # 将pi放入列表
            outputs.append(block_out)

            if i < 2:
                # 对ri进行卷积调整通道数
                route = self.route_blocks_2[i](route)
                # 对ri进行放大，使其尺寸和c_{i+1}保持一致
                route = self.upsample(route)

        return outputs

    def get_pred(self,
                 outputs,
                 im_shape=None,
                 anchors = [10,13,16,30,33,23,30,61,62,45,59,119,116,90,156,198,373,326],
                 anchors_masks = [[6,7,8],[3,4,5],[0,1,2]],
                 valid_thresh = 0.01):
        downsample = 32
        total_boxes = []
        total_scores = []
        for i,out in enumerate(outputs):
            anchor_masks = anchors_masks[i]
            anchors_this_level = []
            for m in anchor_masks:
                anchors_this_level.append(anchors[2 * m])
                anchors_this_level.append(anchors[2 * m + 1])

            boxes, scores = paddle.vision.ops.yolo_box(
                x=out,
                img_size=im_shape,
                anchors=anchors_this_level,
                class_num=self.num_classes,
                conf_thresh=valid_thresh,
                downsample_ratio=downsample,
                name="yolo_box" + str(i)
            )

        total_boxes.append(boxes)
        total_scores.append(paddle.transpose(scores,perm=[0,2,1]))
        downsample = downsample // 2

        yolo_boxes = paddle.concat(total_boxes,axis=1)
        yolo_scores = paddle.concat(total_scores,axis=2)
        return yolo_boxes,yolo_scores